import torch
import torchvision
import torchvision.transforms as transforms
from mymodel import DesNet, ResNet
import os
import matplotlib.pyplot as plt
if __name__ == "__main__":
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    batch_size = 4
    root = '/root/autodl-tmp/dataset/cifar-10-python' if torch.cuda.is_available() else 'E:\\MyBackups\\files\\Datasets\\cifar-10-python'
    device = "cuda" if torch.cuda.is_available() else "cpu"
    testset = torchvision.datasets.CIFAR10(root=root, train=False,
                                           download=False, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                             shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    dataiter = iter(testloader)
    DIR = "DenseNet"
    net = DesNet()
    net.to(device)
    net.load_state_dict(torch.load(os.path.join(DIR, "model.pth")))

    correct = 0
    total = 0
    # since we're not training, we don't need to calculate the gradients for our outputs
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            # calculate outputs by running images through the network
            outputs = net(images)
            # the class with the highest energy is what we choose as prediction
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

    # prepare to count predictions for each class
    correct_pred = {classname: 0 for classname in classes}
    total_pred = {classname: 0 for classname in classes}

    # again no gradients needed
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = net(images)
            _, predictions = torch.max(outputs, 1)
            # collect the correct predictions for each class
            for label, prediction in zip(labels, predictions):
                if label == prediction:
                    correct_pred[classes[label]] += 1
                total_pred[classes[label]] += 1


    # print accuracy for each class
    for classname, correct_count in correct_pred.items():
        accuracy = 100 * float(correct_count) / total_pred[classname]
        print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')

    def plot_curve(x, y, ylabel, title, show=False):
        plt.figure()
        plt.ylabel(ylabel)
        plt.title(title)
        plt.plot(x, y)
        plt.savefig(os.path.join(DIR, title+".png"))
        if show:
            plt.show()


    train_hist = torch.load(os.path.join(DIR, "train_hist.pth"))
    train_acc = train_hist["train_acc"]
    train_loss = train_hist["train_loss"]
    plot_curve(list(range(len(train_acc))), train_acc, "acc", "train acc")
    plot_curve(list(range(len(train_loss))), train_loss, "loss", "train loss")
